#!/usr/bin/env python
# encoding: utf-8
"""支持加密密钥的文件分发系统"""

import os
import time
import hashlib
import threading
import paramiko
from scp import SCPClient, SCPException
from tqdm import tqdm
import concurrent.futures
import logging
from pathlib import Path
from stat import S_ISDIR

# 日志配置
logging.basicConfig(
    filename="distribute.log",
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    filemode="w"
)
logger = logging.getLogger(__name__)


class FileDistributor:
    def __init__(self, hosts, username, key_path=None, key_password=None, password=None):
        """
        这是整个文件分发系统的核心类，包含多个方法用于实现不同的功能。
        增强版初始化参数：
        - key_password: 加密私钥的密码
        """
        self.hosts = hosts
        self.username = username
        self.key_path = Path(key_path) if key_path else None
        self.key_password = key_password
        self.password = password
        self.results = {}
        self.lock = threading.Lock()
        self.connection_pool = {}

    def _get_ssh_client(self, host):
        """支持加密密钥的SSH连接"""
        if host in self.connection_pool:
            return self.connection_pool[host]

        client = paramiko.SSHClient()
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        try:
            if self.key_path and self.key_path.exists():
                # 尝试不同密钥类型
                key = None
                key_classes = [
                    paramiko.RSAKey,
                    paramiko.Ed25519Key,
                    paramiko.ECDSAKey,
                    paramiko.DSSKey
                ]

                for key_class in key_classes:
                    try:
                        key = key_class.from_private_key_file(
                            str(self.key_path),
                            password=self.key_password
                        )
                        break
                    except paramiko.ssh_exception.SSHException:
                        continue

                if not key:
                    raise ValueError("无法识别的密钥类型")

                client.connect(host, username=self.username, pkey=key, timeout=15)
            elif self.password:
                client.connect(host, username=self.username, password=self.password, timeout=15)
            else:
                raise ValueError("必须提供认证方式")

            # 检测操作系统类型
            _, stdout, _ = client.exec_command("uname -s", timeout=5)
            os_type = 'windows' if 'windows' in stdout.read().decode().lower() else 'unix'
            self.connection_pool[host] = (client, os_type)
            return self.connection_pool[host]

        except paramiko.ssh_exception.PasswordRequiredException:
            error_msg = f"[{host}] 密钥需要密码"
            logger.error(error_msg)
            raise Exception(error_msg)
        except Exception as e:
            logger.error(f"[{host}] 连接失败: {str(e)}")
            client.close()
            with self.lock:
                self.results[host] = {"status": "Failed", "error": str(e)}
            return None, None

    def _normalize_path(self, path, os_type):
        """智能路径格式化"""
        if os_type == 'windows':
            return path.replace('/', '\\').rstrip('\\') + '\\'
        return path.replace('\\', '/').rstrip('/') + '/'

    def _md5_checksum(self, file_path):
        """计算文件校验码"""
        hash_md5 = hashlib.md5()
        with open(file_path, "rb") as f:
            for chunk in iter(lambda: f.read(4096), b""):
                hash_md5.update(chunk)
        return hash_md5.hexdigest()

    def _verify_remote_file(self, ssh, remote_path, local_md5, os_type):
        """远程文件校验"""
        remote_path = self._normalize_path(remote_path, os_type)
        cmd = (
            f'CertUtil -hashfile "{remote_path}" MD5' if os_type == 'windows'
            else f'md5sum "{remote_path}"'
        )
        try:
            _, stdout, _ = ssh.exec_command(cmd, timeout=30)
            output = stdout.read().decode().lower()
            remote_md5 = (
                output.split()[1] if os_type == 'windows'
                else output.split()[0]
            )
            return remote_md5 == local_md5
        except Exception as e:
            logger.warning(f"校验失败: {str(e)}")
            return False

    def _transfer_with_retry(self, host, local_path, remote_path, max_retries=3):
        """带重试机制的传输核心"""
        ssh, os_type = self._get_ssh_client(host)
        if not ssh:
            return False

        local_path = Path(local_path)
        remote_path = self._normalize_path(remote_path, os_type)
        remote_full_path = remote_path + local_path.name
        local_md5 = self._md5_checksum(local_path) if local_path.is_file() else None

        for attempt in range(max_retries):
            try:
                # 创建目标目录
                mkdir_cmd = (
                    f'mkdir "{remote_path}"' if os_type == 'windows'
                    else f'mkdir -p "{remote_path}"'
                )
                ssh.exec_command(mkdir_cmd, timeout=10)

                # 带进度传输
                with SCPClient(ssh.get_transport(), socket_timeout=60) as scp:
                    file_size = os.path.getsize(local_path)

                    progress_bar = tqdm(
                        total=file_size,
                        desc=f"{host} 传输进度",
                        unit="B",
                        unit_scale=True,
                        leave=False
                    )

                    def update_progress(transferred, total):
                        progress_bar.update(transferred - progress_bar.n)

                    scp.put(str(local_path), remote_full_path, recursive=True, progress=update_progress)
                    progress_bar.close()

                # 校验文件完整性
                if not self._verify_remote_file(ssh, remote_full_path, local_md5, os_type):
                    raise SCPException("文件校验失败")

                with self.lock:
                    self.results[host] = {"status": "Success", "retries": attempt + 1}
                return True

            except (SCPException, paramiko.SSHException) as e:
                logger.warning(f"[{host}] 第{attempt + 1}次尝试失败: {str(e)}")
                time.sleep(2 ** attempt)  # 指数退避
                continue
            finally:
                if attempt == max_retries - 1:
                    with self.lock:
                        self.results[host] = {"status": "Failed", "error": str(e)}
        return False

    def distribute(self, local_path, remote_path, max_workers=50):
        """执行分发任务"""
        local_path = Path(local_path)
        if not local_path.exists():
            logger.error(f"本地路径不存在: {local_path}")
            raise FileNotFoundError(f"本地文件不存在: {local_path}")

        # 预计算MD5
        local_md5 = self._md5_checksum(local_path) if local_path.is_file() else None

        with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
            futures = {
                executor.submit(
                    self._transfer_with_retry,
                    host,
                    local_path,
                    remote_path
                ): host for host in self.hosts
            }

            with tqdm(total=len(self.hosts), desc="总体进度", unit="host") as main_bar:
                for future in concurrent.futures.as_completed(futures):
                    host = futures[future]
                    try:
                        future.result()
                    except Exception as e:
                        logger.error(f"[{host}] 未捕获异常: {str(e)}")
                    finally:
                        main_bar.update(1)

        # 清理连接池
        for host in self.connection_pool.copy():
            try:
                self.connection_pool[host][0].close()
                del self.connection_pool[host]
            except Exception as e:
                logger.error(f"[{host}] 关闭连接失败: {str(e)}")

        self._generate_report()

    def _generate_report(self):
        """生成详细报告"""
        success = [h for h, r in self.results.items() if r["status"] == "Success"]
        failed = [h for h, r in self.results.items() if r["status"] == "Failed"]

        print("\n" + "=" * 40)
        print(f"{'分发报告':^40}")
        print("=" * 40)
        print(f"▪ 总主机数: {len(self.hosts)}")
        print(f"✓ 成功: {len(success)}")
        print(f"✗ 失败: {len(failed)}")

        if failed:
            print("\n失败详情:")
            for host in failed:
                error = self.results[host].get("error", "未知错误")
                print(f"  {host}: {error}")

        logger.info(f"分发完成 - 成功率: {len(success) / len(self.hosts):.2%}")


if __name__ == "__main__":
    # 使用示例（需要根据实际情况修改）
    hosts = ["113.248.37.162"]

    distributor = FileDistributor(
        hosts=hosts,
        username="Administrator",
        key_path=r"C:\Users\macroh\.ssh\id_rsa",  # 私钥文件路径
        key_password="123456789"  # 密钥密码
    )

    distributor.distribute(
        local_path=r"E:\AAA-project\muchen_project\文件分发系统\text\六百六十六",  # 本地文件路径
        remote_path=r"D:/123456789/",  # 远程目标目录
        max_workers=50  # 并发线程数
    )

"""
scp D:\1122334455  Administrator@212.192.215.208:C:\接收目录\
    
"""